import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
import { logger } from "./logger.js";

const MAX_DATA_POINTS = 60; // This is now irrelevant but kept for structure

app.registerExtension({
    name: "Comfy.lora_train.WorkflowPerformanceMonitorNode",
    
    setup() {
        // Official API-based event listening, per documentation
        // Official API-based event listening, per documentation
     // Official API-based event listening, per documentation
     api.addEventListener("execution_start", () => {
        console.log("📈 [WorkflowMonitor] Event: execution_start. Starting recording.");
        const nodes = app.graph.findNodesByType("WorkflowPerformanceMonitorNode");
        nodes.forEach(node => node.startRecording?.());
    });

    api.addEventListener("status", ({ detail }) => {
        // This event fires when the queue status changes.
        // When the queue is empty, we stop recording.
        // As confirmed by logger, the path is detail.exec_info.
        if (detail?.exec_info?.queue_remaining === 0) {
            console.log("📈 [WorkflowMonitor] Event: status (queue_remaining: 0). Stopping recording.");
            const nodes = app.graph.findNodesByType("WorkflowPerformanceMonitorNode");
            nodes.forEach(node => node.stopRecording?.());
        }
    });
    },

    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name !== "WorkflowPerformanceMonitorNode") {
            return;
        }

        console.log("📈 [WorkflowMonitor] Registering node.");

        nodeType.size = [320, 260];

        const onNodeCreated = nodeType.prototype.onNodeCreated;
        nodeType.prototype.onNodeCreated = function () {
            onNodeCreated?.apply(this, arguments);

            this.performanceData = { cpu: [], ram: [], gpu: [] };
            this.legendHitboxes = {};
            this.visibility = { cpu: true, ram: true, gpu: true };
            this.isRecording = false;
            this.hasExecuted = false;
            
            this.handleDataUpdate = (data) => {
                if (!this.isRecording || !data || !data.cpu || !data.memory || !data.gpu) {
                    return;
                }
                
                this.performanceData.cpu.push(data.cpu.percent);
                this.performanceData.ram.push(data.memory.percent);
                
                const gpuKeys = Object.keys(data.gpu);
                const gpuPercent = gpuKeys.length > 0 ? data.gpu[gpuKeys[0]].percent : 0;
                this.performanceData.gpu.push(gpuPercent);
                
                this.setDirtyCanvas(true, false);
            };
            
            if (window.SystemMonitorInstance) {
                console.log("📈 [WorkflowMonitor] Registering callback to SystemMonitorInstance.");
                window.SystemMonitorInstance.registerDataCallback(this.handleDataUpdate);
            } else {
                 console.warn("📈 [WorkflowMonitor] SystemMonitorInstance not found. Node will not function.");
            }
        };

        const onRemoved = nodeType.prototype.onRemoved;
        nodeType.prototype.onRemoved = function () {
            if (window.SystemMonitorInstance && this.handleDataUpdate) {
                console.log("📈 [WorkflowMonitor] Unregistering callback from SystemMonitorInstance.");
                window.SystemMonitorInstance.unregisterDataCallback(this.handleDataUpdate);
            }
            onRemoved?.apply(this, arguments);
        };
        
        nodeType.prototype.startRecording = function() {
            if (this.isRecording) return; // Already recording
            console.log(`📈 [WorkflowMonitor] Starting recording for node ${this.id}`);
            this.isRecording = true;
            this.hasExecuted = true;
            this.performanceData = { cpu: [], ram: [], gpu: [] };
            this.setDirtyCanvas(true, true); // Force full redraw
        };

        nodeType.prototype.stopRecording = function() {
            if (!this.isRecording) return;
            console.log(`📈 [WorkflowMonitor] Stopping recording for node ${this.id}`);
            this.isRecording = false;
            this.setDirtyCanvas(true, true); // Force full redraw
        };
        
        nodeType.prototype.onMouseDown = function (e, pos) {
            if(this.flags.collapsed) return;

            for (const id in this.legendHitboxes) {
                const box = this.legendHitboxes[id];
                if (pos[0] > box.x && pos[0] < box.x + box.width &&
                    pos[1] > box.y && pos[1] < box.y + box.height) {
                    
                    this.visibility[id] = !this.visibility[id];
                    this.setDirtyCanvas(true, false);
                    return;
                }
            }
        };
        
        nodeType.prototype.onDrawForeground = function (ctx) {
            if (this.flags.collapsed) return;

            const chartRect = { x: 15, y: 30, width: this.size[0] - 30, height: this.size[1] - 70 };
            
            this.drawTitle(ctx);
            this.drawGrid(ctx, chartRect);
            
            if (this.visibility.gpu) this.drawMetric(ctx, this.performanceData.gpu, "#1982c4", chartRect);
            if (this.visibility.ram) this.drawMetric(ctx, this.performanceData.ram, "#ffca3a", chartRect);
            if (this.visibility.cpu) this.drawMetric(ctx, this.performanceData.cpu, "#8ac926", chartRect);

            this.drawLegend(ctx, chartRect);
        };

        nodeType.prototype.drawTitle = function(ctx) {
            ctx.fillStyle = "#fff";
            ctx.font = "bold 14px Arial";
            ctx.textAlign = "center";
            
            let title = "Workflow Performance (Idle)";
            if (this.isRecording) {
                title = "Workflow Performance (Recording...)";
            } else if (this.hasExecuted) {
                title = "Workflow Performance (Finished)";
            }
            ctx.fillText(title, this.size[0] / 2, 20);
        };
        
        nodeType.prototype.drawGrid = function(ctx, rect) {
            ctx.strokeStyle = "#333";
            ctx.lineWidth = 1;
            ctx.font = "10px Arial";
            ctx.fillStyle = "#888";
            ctx.textAlign = "right";

            for (let i = 0; i <= 4; i++) {
                const p = i / 4;
                const y = rect.y + p * rect.height;
                const value = 100 - p * 100;

                ctx.beginPath();
                ctx.moveTo(rect.x, y);
                ctx.lineTo(rect.x + rect.width, y);
                ctx.stroke();
                
                ctx.fillText(`${value}%`, rect.x - 5, y + 3);
            }
            ctx.textAlign = "left";
        };
        
        nodeType.prototype.drawMetric = function(ctx, data, color, rect) {
            if (data.length < 2) return;

            const width = rect.width;
            const height = rect.height;
            const xStep = data.length > 1 ? width / (data.length - 1) : width;
            
            const grad = ctx.createLinearGradient(0, rect.y, 0, rect.y + height);
            grad.addColorStop(0, `${color}60`);
            grad.addColorStop(1, `${color}00`);
            ctx.fillStyle = grad;
            
            ctx.beginPath();
            ctx.moveTo(rect.x, rect.y + height);
            data.forEach((d, i) => {
                const x = rect.x + i * xStep;
                const y = rect.y + height - (d / 100) * height;
                ctx.lineTo(x, y);
            });
            ctx.lineTo(rect.x + (data.length-1) * xStep, rect.y + height);
            ctx.closePath();
            ctx.fill();

            ctx.strokeStyle = color;
            ctx.lineWidth = 2;
            ctx.beginPath();
            data.forEach((d, i) => {
                const x = rect.x + i * xStep;
                const y = rect.y + height - (d / 100) * height;
                if (i === 0) ctx.moveTo(x, y);
                else ctx.lineTo(x, y);
            });
            ctx.stroke();
        };

        nodeType.prototype.drawLegend = function(ctx, rect) {
            const legendY = rect.y + rect.height + 25;
            const legendItems = [
                { id: "gpu", label: "GPU", color: "#1982c4", value: this.performanceData.gpu.at(-1) || 0 },
                { id: "ram", label: "RAM", color: "#ffca3a", value: this.performanceData.ram.at(-1) || 0 },
                { id: "cpu", label: "CPU", color: "#8ac926", value: this.performanceData.cpu.at(-1) || 0 }
            ];

            ctx.font = "12px Arial";
            let currentX = rect.x;
            this.legendHitboxes = {};

            for (const item of legendItems) {
                const isVisible = this.visibility[item.id];
                ctx.globalAlpha = isVisible ? 1.0 : 0.4;

                ctx.fillStyle = item.color;
                ctx.fillRect(currentX, legendY - 8, 10, 10);
                
                ctx.fillStyle = "#fff";
                const text = `${item.label}: ${item.value.toFixed(1)}%`;
                ctx.fillText(text, currentX + 15, legendY);
                
                const itemWidth = ctx.measureText(text).width + 20;
                this.legendHitboxes[item.id] = {
                    x: currentX,
                    y: legendY - 10,
                    width: itemWidth,
                    height: 20
                };
                
                currentX += itemWidth + 15;
            }
            ctx.globalAlpha = 1.0;
        };
    }
});